package main

import (
	"fmt"
	"net"
	"shell/common"
	"shell/common/packages"
	"shell/global"
	"shell/utils"
)

/**
	处理目标节点是公网代理端的请求
	- shell命令
	- socks5转发
	- 内网端口转发
**/

// 处理shell命令请求，如果发生错误，就关闭stream流和连接
func HandleShellCmdReqToProxy(req *packages.ShellCmdReq, stream *common.QuicStream) error {
	resp := new(packages.ShellCmdResp)
	resp.FromReq(req, Role)
	output, code, err := utils.RunCmd(string(req.ShellCmd))
	resp.CmdExitCode = code
	if err != nil {
		resp.CmdOutput = []byte(err.Error())
	} else {
		resp.CmdOutput = output
	}
	err = stream.WritePackage(resp)
	if err != nil {
		stream.CloseConn(common.CloseByProcess)
		return err
	}
	return nil
}

// 处理socks5转发
func HandleSingleSocks5StreamToProxy(stream *common.QuicStream, req *packages.Socks5RedirReq) {
	resp := req.ToResp(Role)
	respFunc := func(errInfo string) {
		resp.IsOk = false
		resp.ErrorInfo = []byte(errInfo)
		err := stream.WritePackage(resp)
		if err != nil {
			common.HandleErr(err)
			return
		}
		stream.WaitForClose(Role, false)
	}
	tcpConn, err := net.Dial("tcp", net.JoinHostPort(string(req.TargetAddr), fmt.Sprintf("%d", req.TargetPort)))
	if err != nil {
		respFunc("tcp dial error: " + err.Error())
		return
	}
	resp.IsOk = true
	err = stream.WritePackage(resp)
	if err != nil {
		_ = tcpConn.Close()
		common.HandleErr(err)
		return
	}
	common.NewTimeoutRWCPairWarp(tcpConn, stream, global.DefaultIdleTime, timeoutWatcher).RedirSyncNoKey()
}

// 处理目的节点是代理端的私有端口转发，阻塞
func HandlePrivatePortRedirStreamToProxy(stream *common.QuicStream, req *packages.PortRedirReq) {
	resp := req.ToResp(packages.RoleProxy)
	tcpConn, err := net.Dial("tcp", net.JoinHostPort(req.RedirDestIP.String(), fmt.Sprintf("%d", req.RedirDestPort)))
	if err != nil {
		resp.IsOk = false
		resp.ErrorInfo = []byte(err.Error())
		err = stream.WritePackage(resp)
		if err != nil {
			common.HandleErr(err)
			return
		}
		stream.WaitForClose(Role, false)
	}
	resp.IsOk = true
	err = stream.WritePackage(resp)
	if err != nil {
		common.HandleErr(err)
		return
	}
	common.NewTimeoutRWCPairWarp(tcpConn, stream, global.DefaultIdleTime, timeoutWatcher).RedirSyncNoKey()
}
